package com.anji.hyperneat.onlinereinforcement;

import com.anji.hyperneat.nd.NDFloatArray;
import com.anji.hyperneat.nd.Util;
import com.anji.hyperneat.nd.NDFloatArray.MatrixIterator;

public class ActivatorNdTdZero extends OnlineLearningActivatorNDBase {

	protected NDFloatArray[] lastGradients;
	protected NDFloatArray[] gradients;
	protected NDFloatArray lastOuts;
	protected NDFloatArray outs;
	protected GridNetTD workingTd;
	protected GridNetTD lastSnapshot;
	
	protected static final float alpha = 0.5f; 
	
	// w = w + alpha * (r + current - last) * last_gradient
	
	public ActivatorNdTdZero(ActivatorNDLR phenotype) {
		super(phenotype);
		workingTd = (GridNetTD) working;
		lastGradients = null;
	}

	@Override
	public NDFloatArray next(NDFloatArray stimuli) {
//		lastSnapshot = workingTd.clone();
		
		
		if (outs != null) lastOuts = outs;
		outs = super.next(stimuli);
		
		if (gradients != null) lastGradients = gradients;
		gradients = calcGradients();
		
		return lastOuts;
	}
	
	@Override
	public void updateNet(NDFloatArray rewards) {
		updateNet(rewards, 1.0f);
	}
	
	@Override
	public void updateNet(NDFloatArray rewards, float learningRateModifier) {
		int[] coords;
		
		if (null != lastGradients) {
			NDFloatArray tdErrs = calcTDErrs( rewards );
			
			for (int layer = 0; layer < working.getWeights().length; layer++) {
			
				// For each weight
				for (MatrixIterator w = working.getWeights()[layer].iterator(); w.hasNext(); w.next()) {
					int[] wCoords = w.getCurrentCoordinates();
//					float val = alpha * sumTdErrorsAndGradients(layer, wCoords, tdErrs);		
//					w.set(w.get() + val);
				}
			}
		}
		working.decayLearningRates();
	}


	/**
	 * @param desiredOuts
	 * @param d
	 */
	protected NDFloatArray[] calcGradients() {
		NDFloatArray[] d = new NDFloatArray[working.getNumLayers()-1];
		int[] coords;
		
		// Output layer first
		int layer = d.length-1; 
		d[layer] = new NDFloatArray(working.getOutputDimension());
		for (MatrixIterator it = working.getOutputs().iterator(); 
			it.hasNext();
			it.next()
		) {
			coords = it.getCurrentCoordinates();

			float val = it.get();
			val = val * (1 - val);
			d[layer].set( val, coords );
		}
		
		// Each layer prior
		for (layer = d.length - 1; layer > 0; layer--) {
			d[layer-1] = new NDFloatArray(working.getLayer(layer).getDimensions());
			for (MatrixIterator it = working.getLayer(layer).iterator(); 
				it.hasNext();
				it.next()
			) {
				coords = it.getCurrentCoordinates();
				float val = it.get();
				val = val * (1 - val) * calcWeightErrSum(coords, working.getWeights()[layer], d[layer]);
				d[layer-1].set( val, coords );
			}
		}
		
		return d;
	}


	

	private float calcWeightErrSum(int[] srcCoords, NDFloatArray wgts, NDFloatArray discrepancyErrs) {
		float sum = 0;
		for (MatrixIterator it = discrepancyErrs.iterator(); 
			it.hasNext();
			it.next()
		) {
			int[] tgtCoords = it.getCurrentCoordinates();
			sum += wgts.get(tgtCoords, srcCoords) * it.get();
		}
		
//		for (MatrixIterator it = wgts.iterator(); 
//		it.hasNext();
//		it.next()
//		) {
//			
//			int[] wgtSrcCoords = getSrcCoordsFromCoordsSet(it.getCurrentCoordinates());
//			
//		}
		
		return sum;
	}
	
	/*
	@Override
	public void updateNet(NDFloatArray rewards) {
		
		NDFloatArray deltas[] = calcDeltas();
		NDFloatArray gradients[] = calcGradients( deltas );
		
		if (null != lastGradients) {
			NDFloatArray tdErrs = calcTDErrs( rewards );
			
			for (int layer = 0; layer < working.getWeights().length; layer++) {
			
				// For each weight
				for (MatrixIterator w = working.getWeights()[layer].iterator(); w.hasNext(); w.next()) {
					int[] wCoords = w.getCurrentCoordinates();
					float val = alpha * sumTdErrorsAndGradients(layer, wCoords, tdErrs);		
					w.set(w.get() + val);
				}
			}
		}
		
		lastGradients = gradients;
	}
	*/

	/*
	 * Calculate the delta terms for gradient descent for each layer;
	 * for Temporal Difference, use gradient of output with respect to weights.
	 * @return The NDFloatArray's of delta terms.
	 */
	
//	protected NDFloatArray[] calcDeltas() {
//		int depth = working.getNumLayers();
//		NDFloatArray[] deltas = new NDFloatArray[depth - 1];	// don't need deltas for input layer
//		
//		for (int layer = depth - 1; layer > 0; layer--) {	// start with the last layer, work backward
//			// deltas is one less than depth, offset one back from the current layer
//			
//			if (layer == depth - 1) {
//				deltas[layer-1] = new NDFloatArray(working.getOutputDimension());	// If it's the output layer just use the output layer dimensions 	
//
//				for (MatrixIterator delta = deltas[layer-1].iterator(); delta.hasNext(); delta.next()) {
//					float deltaValue = working.getLayer(layer).getFromRawCoordinate( delta.getRawCoordinate() );	// For output deltas, use complete set of coords
//					deltaValue = deltaValue * (1 - deltaValue);
//					delta.set(deltaValue);
//				}
//			} else {
//				/* Otherwise, it's the dimensions of the output layer * the dimensions of the current layer.
//				 * The reason is, that you calculate the delta for the previous layers by
//				 * calculating the delta for each pair of output layer nodes and current layer nodes.
//				 * Look at calcDeltaSum, it will make more sense.
//				 */	
//				deltas[layer-1] = new NDFloatArray(working.getOutputDimension(), working.getLayer(layer).getDimensions());
//					
//				for (MatrixIterator delta = deltas[layer-1].iterator(); delta.hasNext(); delta.next()) {
//					/* for non-output deltas, extract the source layer coordinates from the delta to get the node for the "output" calculation
//					 * 
//					 */
//					int[] deltaCoords = delta.getCurrentCoordinates();
//					int[] srcCoords = Util.getSrcCoordsFromCoordsSet(deltaCoords);
//					int[] tgtCoords = Util.getTgtCoordsFromCoordsSet(deltaCoords);
//					float deltaValue = working.getLayer(layer).get(srcCoords);	 
////					deltaValue = (deltaValue * (1 - deltaValue)) * calcDeltaSum(layer, deltas, delta);
//					deltaValue = (deltaValue * (1 - deltaValue));
//					// This will only work for 3 layer nets.  But who cares.
//					deltaValue = deltaValue * deltas[layer].get(tgtCoords) * working.getWeights()[layer].get(tgtCoords, srcCoords);
//					delta.set(deltaValue);
//				}
//			}
//		}
//		
//		return deltas;
//	}


//	protected NDFloatArray[] calcGradients(NDFloatArray[] deltas) {
//		int depth = working.getNumLayers();
//		NDFloatArray[] gradients = new NDFloatArray[depth - 1];
//		for (int layer = 0; layer < depth - 1; layer++) {
//			
//			gradients[layer] = new NDFloatArray(
//					// deltas will be 0 thru depth-2 representing layers 1 thru depth-1
//					deltas[layer].getDimensions()
//					// but need the incoming activation value, hence the same numerical layer is used
//					, working.getLayer(layer).getDimensions());
//			
//			for (MatrixIterator delta = deltas[layer].iterator(); delta.hasNext(); delta.next()) {
//				int[] deltaCoords = delta.getCurrentCoordinates();
//				for (MatrixIterator incomingActivation = working.getLayer(layer).iterator()
//						; incomingActivation.hasNext()
//						; incomingActivation.next()) {
//					float value = delta.get() * incomingActivation.get();
//					gradients[layer].set(value, Util.joinCoords(deltaCoords, incomingActivation.getCurrentCoordinates()));
//				}
//			}
//		}		
//		return gradients;
//	}
	
	protected NDFloatArray calcTDErrs(NDFloatArray rewards) {
		NDFloatArray tdErrs = new NDFloatArray(working.getOutputs().getDimensions());
		
		for (MatrixIterator tdErr = tdErrs.iterator(); tdErr.hasNext(); tdErr.next()) {
			int[] tdErrCoords = tdErr.getCurrentCoordinates();
			int tdRawCoord = tdErr.getRawCoordinate();
			float value = rewards.get(tdErrCoords)
				// gammas should have one additional layer beyond weights, for the output
				+ workingTd.getOutputGamma(tdErrCoords) 
				* working.getOutputs().getFromRawCoordinate(tdRawCoord)
				- lastOuts.getFromRawCoordinate(tdRawCoord);
			
			tdErr.set(value);
		}
		
		return tdErrs;
	}

//	protected float sumTdErrorsAndGradients(int layer, int[] wCoords, NDFloatArray tdErrs) {
//		float sum = 0;
//		if (layer == working.getNumLayers()-2) {
//			sum += tdErrs.get(Util.getTgtCoordsFromCoordsSet(wCoords)) * lastGradients[layer].get(wCoords);
//		} else {
//			for (MatrixIterator err = tdErrs.iterator(); err.hasNext(); err.next()) {
//				sum += err.get() * lastGradients[layer].get(wCoords, err.getCurrentCoordinates());
//			}
//		}
//		
//		return sum;
//	}

	
}
